package db

import (
	"context"
	"fmt"
	"github.com/zeromicro/go-zero/core/logx"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"payme3000/framework/errorx"
	"sync"
)

var (
	mongoDbLock sync.RWMutex
	mongoMap    = map[Type]*mongo.Client{}
	mongoDbMap  = map[Type]*mongo.Database{}
)

func GetMongoDbUri(config *Config) string {
	if config.ReplicaSet != "" {
		if config.ReadPreference != "primary" {
			config.ReadPreference = "secondaryPreferred"
		}

		return fmt.Sprintf(
			"mongodb://%s:%s@%s/?replicaSet=%s&readPreference=%s&retryWrites=false&ssl=%t&authSource=admin&w=majority",
			config.Username,
			config.Password,
			config.Host,
			config.ReplicaSet,
			config.ReadPreference,
			config.IsSSL,
		)
	}

	return fmt.Sprintf(
		"mongodb://%s:%s@%s/?connect=direct&ssl=%t",
		config.Username,
		config.Password,
		config.Host,
		config.IsSSL,
	)
}

func GetMongodb(dbType Type) *mongo.Database {
	mongoDbLock.RLock()
	defer mongoDbLock.RUnlock()

	return mongoDbMap[dbType]
}

func CreateMongoDb(dbType Type, config *Config) error {
	mongoDbLock.Lock()
	defer mongoDbLock.Unlock()

	// 锁定后再检查一次
	if client := mongoMap[dbType]; client != nil {
		return nil
	}

	client, err := mongo.Connect(
		context.Background(),
		options.Client().ApplyURI(GetMongoDbUri(config)),
	)
	if err != nil {
		return errorx.Throw(context.Background(), err)
	}

	mongoMap[dbType] = client
	mongoDbMap[dbType] = client.Database(config.Db)

	logx.Infof("create mongodb done with dbType:%d, host:%s, db:%s", dbType, config.Host, config.Db)

	return nil
}

func StopMongoDb() {
	mongoDbLock.RLock()
	defer mongoDbLock.RUnlock()

	for _, mongo := range mongoMap {
		errorx.Catch(context.Background(), mongo.Disconnect(context.Background()))
	}
}
